using System.Collections.Generic;

namespace System.Linq
{
	/// <summary>
	/// Extension methods for Linq to Trees
	/// </summary>
	internal static class EnumerableExtensions
	{
		/// <summary>
		/// Traverses the specified source.
		/// </summary>
		/// <typeparam name="T"></typeparam>
		/// <param name="source">The source.</param>
		/// <param name="predicate">The predicate.</param>
		/// <returns></returns>
		public static IEnumerable<T> Traverse<T>(this IEnumerable<T> source, Func<T, bool> predicate) where T : IEnumerable<T>
		{
			List<T> list = new List<T>();
			Y<IEnumerable<T>>(f => items =>
			{
				list.AddRange(items.Where(predicate));
				foreach (T i in items)
				{
					f(i);
				}
			})(source);
			return list;
		}

		/// <summary>
		/// Traverses the specified source.
		/// </summary>
		/// <typeparam name="T"></typeparam>
		/// <param name="source">The source.</param>
		/// <param name="getChildren">The get children.</param>
		/// <param name="predicate">The predicate.</param>
		/// <returns></returns>
		public static IEnumerable<T> Traverse<T>(this T source, Func<T, IEnumerable<T>> getChildren, Func<T, bool> predicate = null)
		{
			if (predicate == null)
			{
				predicate = x => true;
			}
			List<T> list = new List<T>();
			if (predicate(source))
			{
				list.Add(source);
			}
			Y<IEnumerable<T>>(f => items =>
			{
				list.AddRange(items.Where(predicate));
				foreach (T i in items)
				{
					f(getChildren(i));
				}
			})(getChildren(source));
			return list;
		}

		/// <summary>
		/// Traverses the specified source.
		/// </summary>
		/// <typeparam name="T"></typeparam>
		/// <param name="source">The source.</param>
		/// <param name="getChildren">The get children.</param>
		/// <param name="predicate">The predicate.</param>
		/// <returns></returns>
		public static IEnumerable<T> Traverse<T>(this IEnumerable<T> source, Func<T, IEnumerable<T>> getChildren, Func<T, bool> predicate = null)
		{
			if (predicate == null)
			{
				predicate = x => true;
			}
			List<T> list = new List<T>();
			Y<IEnumerable<T>>(f => items =>
			{
				list.AddRange(items.Where(predicate));
				foreach (T i in items)
				{
					f(getChildren(i));
				}
			})(source);
			return list;
		}

		/// <summary>
		/// Traverses the specified source.
		/// </summary>
		/// <typeparam name="TSource">The type of the source.</typeparam>
		/// <typeparam name="TResult">The type of the result.</typeparam>
		/// <param name="source">The source.</param>
		/// <param name="getChildren">The get children.</param>
		/// <param name="transform">The transform.</param>
		/// <param name="predicate">The predicate.</param>
		/// <returns></returns>
		public static IEnumerable<TResult> Traverse<TSource, TResult>(this TSource source,
		                                                                Func<TSource, IEnumerable<TSource>> getChildren,
		                                                                Func<TSource, IEnumerable<TResult>> transform,
		                                                                Func<TResult, bool> predicate = null)
		{
			if (predicate == null)
			{
				predicate = x => true;
			}
			List<TResult> list = new List<TResult>();
			list.AddRange(transform(source).Where(predicate));
			Y<IEnumerable<TSource>>(f => items =>
			{
				foreach (TSource i in items)
				{
					list.AddRange(transform(i).Where(predicate));
					f(getChildren(i));
				}
			})(getChildren(source));
			return list;
		}

		private static Action<T> Y<T>(Func<Action<T>, Action<T>> f)
		{
			return a => f(Y(f))(a);
		}
	}
}